import cv2
import numpy as np

class yolov5_face_detector():
    def __init__(self, yolo_type, confThreshold=0.5, nmsThreshold=0.5, objThreshold=0.5):
        anchors = [[4, 5, 8, 10, 13, 16], [23, 29, 43, 55, 73, 105], [146, 217, 231, 300, 335, 433]]
        num_classes = 1
        self.nl = len(anchors)  # 预测特征图数
        self.na = len(anchors[0]) // 2  # 预测特征图上每个位置生成anchor数
        self.no = num_classes + 5 + 10  # yolo网络在预测特征图每个位置输出no个数，预测类别、bbox回归参数、关键点回归参数
        self.grid = [np.zeros(1)] * self.nl
        self.stride = np.array([8., 16., 32.]) # 三个预测特征图的下采样倍数
        self.anchor_grid = np.asarray(anchors, dtype=np.float32).reshape(self.nl, -1, 2)
        self.inpWidth = 640  # yolo网络输入图片宽
        self.inpHeight = 640  # yolo网络输入图片高
        self.net = cv2.dnn.readNet("utils/yolov5_face_weights/" + yolo_type + '-face.onnx')
        self.confThreshold = confThreshold
        self.nmsThreshold = nmsThreshold
        self.objThreshold = objThreshold

    def _make_grid(self, nx=20, ny=20):
        xv, yv = np.meshgrid(np.arange(ny), np.arange(nx))
        return np.stack((xv, yv), 2).reshape((-1, 2)).astype(np.float32)

    def detect(self, srcimg):
        # Input Image pre-transform
        blob = cv2.dnn.blobFromImage(srcimg, 1 / 255.0, (self.inpWidth, self.inpHeight), [0, 0, 0], swapRB=True,
                                     crop=False)
            # 图片预处理，像素值归一化 --> resize到网络输入尺寸 --> 交换RB维度(cv2.imread读入默认为bgr)
        # Sets the input to the network
        self.net.setInput(blob)

        # Runs the forward pass to get output of the output layers
        outs = self.net.forward(self.net.getUnconnectedOutLayersNames())[0]
            # outs输出所有预测bbox的参数
            # 每个bbox的参数及顺序为：4个anchor回归参数 + 1个obj conf预测参数 + 10个landmark回归参数 + 1个类别预测参数 

        # inference output
        outs[..., [0, 1, 2, 3, 4, 15]] = 1 / (1 + np.exp(-outs[..., [0, 1, 2, 3, 4, 15]]))  ###sigmoid
            # bbox回归、obj conf、类别预测参数输出后需进行sigmoid才可得最终输出
        row_ind = 0
        for i in range(self.nl):
            # 循环对三个预测特征图的输出进行处理
            h, w = int(self.inpHeight / self.stride[i]), int(self.inpWidth / self.stride[i])
            # h,w 当前预测特征图的尺寸
            length = int(self.na * h * w)  # length 当前预测特征图的anchor数
            if self.grid[i].shape[2:4] != (h, w):
                self.grid[i] = self._make_grid(w, h)

            g_i = np.tile(self.grid[i], (self.na, 1))
            a_g_i = np.repeat(self.anchor_grid[i], h * w, axis=0)
            outs[row_ind:row_ind + length, 0:2] = (outs[row_ind:row_ind + length, 0:2] * 2. - 0.5 + g_i) * int(
                self.stride[i])
            outs[row_ind:row_ind + length, 2:4] = (outs[row_ind:row_ind + length, 2:4] * 2) ** 2 * a_g_i

            outs[row_ind:row_ind + length, 5:7] = outs[row_ind:row_ind + length, 5:7] * a_g_i + g_i * int(
                self.stride[i])  # landmark x1 y1
            outs[row_ind:row_ind + length, 7:9] = outs[row_ind:row_ind + length, 7:9] * a_g_i + g_i * int(
                self.stride[i])  # landmark x2 y2
            outs[row_ind:row_ind + length, 9:11] = outs[row_ind:row_ind + length, 9:11] * a_g_i + g_i * int(
                self.stride[i])  # landmark x3 y3
            outs[row_ind:row_ind + length, 11:13] = outs[row_ind:row_ind + length, 11:13] * a_g_i + g_i * int(
                self.stride[i])  # landmark x4 y4
            outs[row_ind:row_ind + length, 13:15] = outs[row_ind:row_ind + length, 13:15] * a_g_i + g_i * int(
                self.stride[i])  # landmark x5 y5
            row_ind += length
        return outs
    
    def do_mosaic(self, img, box, mosaic_level):
        """
        :param rgb_img
        :param int x :  马赛克左顶点
        :param int y:  马赛克左顶点
        :param int w:  马赛克宽
        :param int h:  马赛克高
        :param int neighbor:  马赛克每一块的宽
        """
        w = int(box[2])
        h = int(box[3])
        x = int(box[0])
        y = int(box[1])
        neighbor = max(int(mosaic_level * 0.2 * w), 3)
        for i in range(0, h, neighbor):
            for j in range(0, w, neighbor):
                # rect = [j + x, i + y]
                x1 = j + x
                y1 = i + y
                color = img[y1][x1].tolist()  # 关键点1 tolist
                left_up = (x1, y1)
                x2 = x1 + neighbor - 1  # 关键点2 减去一个像素
                y2 = y1 + neighbor - 1
                if x2 > x + w:
                    x2 = x + w
                if y2 > y + h:
                    y2 = y + h
                right_down = (x2, y2)
                cv2.rectangle(img, left_up, right_down, color, -1)

        return img

    def drawPred(self, frame, conf, left, top, right, bottom, landmark):
        # Draw a bounding box.
        cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255), thickness=2)
        # label = '%.2f' % conf
        # Display the label at the top of the bounding box
        # labelSize, baseLine = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
        # top = max(top, labelSize[1])
        # cv2.putText(frame, label, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), thickness=2)
        for i in range(5):
            cv2.circle(frame, (landmark[i * 2], landmark[i * 2 + 1]), 1, (0, 255, 0), thickness=-1)
        return frame
    
    def postprocess(self, frame, outs, mark, mosaic_level=0.6):
        frameHeight = frame.shape[0]
        frameWidth = frame.shape[1]
        ratioh, ratiow = frameHeight / self.inpHeight, frameWidth / self.inpWidth
        # Scan through all the bounding boxes output from the network and keep only the
        # ones with high confidence scores. Assign the box's class label as the class with the highest score.

        confidences = []
        boxes = []
        landmarks = []
        for detection in outs:
            confidence = detection[15]
            objectness = detection[4]
            if confidence > self.confThreshold and objectness > self.objThreshold:
            # if detection[4] > self.objThreshold:
                center_x = int(detection[0] * ratiow)
                center_y = int(detection[1] * ratioh)
                width = int(detection[2] * ratiow)
                height = int(detection[3] * ratioh)
                left = int(center_x - width / 2)
                top = int(center_y - height / 2)

                confidences.append(float(confidence))
                boxes.append([left, top, width, height])
                landmark = detection[5:15] * np.tile(np.float32([ratiow, ratioh]), 5)
                landmarks.append(landmark.astype(np.int32))
        # Perform non maximum suppression to eliminate redundant overlapping boxes with
        # lower confidences.
        indices = cv2.dnn.NMSBoxes(boxes, confidences, self.confThreshold, self.nmsThreshold)
        box_num = 0
        real_boxes = []
        for i in indices:
            box_num += 1  # 检测到的目标数量
            # i = i[0]
            box = boxes[i]
            left = box[0]
            top = box[1]
            width = box[2]
            height = box[3]
            right = left + width
            bottom = top - height
            real_boxes.append([left, top, right, bottom])
            landmark = landmarks[i]
            if mark == 'mosaic':
                frame = self.do_mosaic(frame, box, mosaic_level)
            elif mark == 'origin':
                frame = self.drawPred(frame, confidences[i], left, top, left + width, top + height, landmark)

        return frame, box_num, real_boxes


if __name__ == "__main__":
    image = cv2.imread('/home/igs/Codefield/Python/Face_mosaic/inputs_face/1.jpeg')
    detector = yolov5_face_detector(yolo_type='yolov5s')
    outs = detector.detect(image)
    image_result, bbox_num = detector.postprocess(image, outs, mark='origin')
